# TODO 1.导包
import os
import pandas as pd
import datetime

from sklearn.preprocessing import LabelEncoder
from sklearn.utils import compute_sample_weight
from xgboost import XGBRegressor, XGBClassifier
from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import mean_squared_error, mean_absolute_error, roc_auc_score
import joblib

import matplotlib.pyplot as plt
import matplotlib

matplotlib.rcParams['axes.unicode_minus']= False #解决负号'-’显示为方块的问题
matplotlib.rcParams['font.family']='Kaiti SC'#可以替换为其他字体

# TODO 2.数据分析 (产生符合分析图保存到data/fig目录)
def ana_data(data):
    #data.info()
    #print(data.head())
    #print(data.columns)

    fig = plt.figure(figsize=(20, 32))

    # 不同年龄的离职率和未离职率
    ax1 = fig.add_subplot(411)
    attrited = data[data['Attrition'] == 1]['Age'].value_counts().sort_index()
    not_attrited = data[data['Attrition'] == 0]['Age'].value_counts().sort_index()

    ax1.plot(attrited.index, attrited.values, label='离职')
    ax1.plot(not_attrited.index, not_attrited.values, label='未离职')

    ax1.set_title('不同年龄的离职率和未离职率')
    ax1.set_xlabel('年龄')
    ax1.set_ylabel('人数')
    ax1.legend()

    plt.tight_layout()


    plt.savefig('../data/fig/age_attrition.png')
    #不同工作满意度的离职率
    ax2 = fig.add_subplot(412)
    attrited = data[data['Attrition'] == 1]['JobSatisfaction'].value_counts().sort_index()
    not_attrited = data[data['Attrition'] == 0]['JobSatisfaction'].value_counts().sort_index()
    ax2.plot(attrited.index, attrited.values, label='离职')
    ax2.plot(not_attrited.index, not_attrited.values, label='未离职')
    ax2.set_title('不同工作满意度的离职率')
    ax2.set_xlabel('工作满意度')
    ax2.set_ylabel('人数')
    ax2.legend()
    plt.tight_layout()
    plt.show()

#TODO 特征工程
def feature_engineering(data):
    data = data.copy(deep=True)
    # 1.删除无用列
    data = data.drop(['Over18', 'StandardHours', 'EmployeeNumber','TrainingTimesLastYear','RelationshipSatisfaction', 'PercentSalaryHike','EducationField'], axis=1)
    # 2.使用LabelEncoder修改JobRole、MaritalStatus的值
    data['JobRole'] = LabelEncoder().fit_transform(data['JobRole'])
    data['MaritalStatus'] = LabelEncoder().fit_transform(data['MaritalStatus'])
    # 3.采用mapping映射方法
    # 手动定义映射关系,修改BusinessTravel的值
    mapping1 = {'Non-Travel': 0, 'Travel_Rarely': 1, 'Travel_Frequently': 2}
    # 使用map()替换
    data['BusinessTravel'] = data['BusinessTravel'].map(mapping1)
    # 手动定义映射关系,修改Department的值
    mapping2 = {'Human Resources': 1, 'Research & Development': 2, 'Sales': 3}
    data['Department'] = data['Department'].map(mapping2)
    # 4.采用热编码处理OverTime
    data = pd.get_dummies(data, columns=['OverTime', 'Gender'], drop_first=True)

    return data

#TODO 数据集划分
def train_test_sp(data):
    x = data.iloc[:, 1:]
    y = data.iloc[:, 0]
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2,
                                                        random_state=88, stratify=y)
    return x_train, y_train, x_test, y_test

def train_grid_search(x_train, y_train):
    param_dict = {
        'n_estimators': [505, 615, 370, 605, 857, 272,635,357,785,342,764,456],
        'max_depth': [1, 2 ,3 ,4],
        'learning_rate': [0.23, 0.33, 0.1, 0.18,0.01,0.5,0.8,0.26]
    }
    cls_weight = compute_sample_weight('balanced', y_train)
    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=88)
    gsc = GridSearchCV(estimator=XGBClassifier(), param_grid=param_dict, cv=skf, scoring='roc_auc')
    gsc.fit(x_train, y_train, sample_weight=cls_weight)
    print(f"最佳参数best_params:{gsc.best_params_}")
    print(f"最佳评分best_score:{gsc.best_score_}")

def model_train(x_train, y_train, x_test, y_test):
    cls_weight = compute_sample_weight('balanced', y_train)
    xgb_model = XGBClassifier(n_estimators=357, max_depth=1, learning_rate=0.1)
    # xgb_model = XGBClassifier(n_estimators=505, max_depth=1, learning_rate=0.1)
    xgb_model.fit(x_train, y_train, sample_weight=cls_weight)
    y_pre = xgb_model.predict_proba(x_test)[:, 1]
    print(f'roc_auc_score:{roc_auc_score(y_test, y_pre)}')
    joblib.dump(xgb_model, '../model/xgb.pkl')

if __name__ == '__main__':
    data = pd.read_csv('../data/train.csv')
    #ana_data(data)
    data = feature_engineering(data)
    x_train, y_train, x_test, y_test = train_test_sp(data)
    #train_grid_search(x_train, y_train)
    model_train(x_train, y_train, x_test, y_test)


